tree grammar GraphrWalker;

options {
  tokenVocab=Graphr;
  ASTLabelType=CommonTree;
  backtrack=true;
}

@header {
  import java.util.Map;
  import java.util.Stack;
  import java.util.HashMap;
  import java.io.BufferedReader;
  import java.io.BufferedWriter;
  import java.io.FileReader;
  import java.io.FileWriter;
  import java.io.FileInputStream;
  import java.io.IOException;
  import org.antlr.runtime.*;
  import org.antlr.runtime.tree.*;
  import org.antlr.stringtemplate.*;
}

@members {
  /** Map variable name to Integer object holding value */
  Map<String, GraphrDataType> memory = new HashMap<String, GraphrDataType>();

  /** Points to functions tracked by tree builder / parser. */
  HashMap<String, CommonTree> functions;
  // points to the beginning function of function markers in input stream
  HashMap<String, Integer> functionMarkers = new HashMap<String, Integer>(); 
  // points to end of function markers in input stream
  HashMap<String, Integer> functionEndMarkers = new HashMap<String, Integer>();
  // a stack of function return values
  Stack<GraphrDataType> functionReturnValues = new Stack<GraphrDataType>();
  // stack for marking where to return to in the input stream when we return from a function
  Stack<Integer> calledFunctionEndMarkers = new Stack<Integer>();

  /** Stack of value scopes.  Stack<Map>; each map holds set of
   *  name/value pairs.
   */
  Stack<Map> mainMemoryScopes = new Stack<Map>();
  Stack<Stack<Map>> functionScopes = new Stack<Stack<Map>>();

  /** Get value of name. If we're in a function, we look at the top function scope in
   *  the functionScopes. Otherwise, look in the mainMemoryScopes.
   *  if we don't find a value, print an error.
  **/
  public GraphrDataType getValue(String name) {
    GraphrDataType val = null;
    if (!functionScopes.isEmpty()) {
      val = lookForValueInStack(functionScopes.peek(), name);
    }
    else {
      val = lookForValueInStack(mainMemoryScopes, name);
    }
    if (val == null) {
      System.out.println("There was no variable found named: " + name);
      System.exit(1);
    }
    return val;
  }
  
  public GraphrDataType lookForValueInStack(Stack<Map> scopes, String name) {
    if (scopes.isEmpty()) {
      return null;
    }
    Map<String, GraphrDataType> values = scopes.pop();
    GraphrDataType val = values.get(name);
    if (val == null) {
      val = lookForValueInStack(scopes, name);
    }
    scopes.push(values);
    return val;
  }
  
  public boolean findAndSetValue(Stack<Map> scopes, String name, GraphrDataType val) {
    if (scopes.isEmpty()) {
      return false;
    }
    boolean foundValue = false;
    Map<String, GraphrDataType> values = scopes.pop();
    if (values.get(name) != null) {
      values.put(name, val);
      foundValue = true;
    }
    else {
      foundValue = findAndSetValue(scopes, name, val);
    }
    scopes.push(values);
    return foundValue;
  }
  
  public GraphrDataType setValue(String name, GraphrDataType val) {
    if (!functionScopes.isEmpty()) {
      boolean set = findAndSetValue(functionScopes.peek(), name, val);
      if (!set) {
        functionScopes.peek().peek().put(name, val);
      }
    }
    else {
      boolean set = findAndSetValue(mainMemoryScopes, name, val);
      if (!set) {
        mainMemoryScopes.peek().put(name, val);
      }
    }
    return val;
  }
  
  public HashMap<String, GraphrDataType> pushScope() {
    HashMap<String, GraphrDataType> newScope = new HashMap<String, GraphrDataType>();
    if (functionScopes.isEmpty()) {
      mainMemoryScopes.push(newScope);
    }
    else {
      functionScopes.peek().push(newScope);
    }
    return newScope;
  }
  
  public void popScope() {
    if (functionScopes.isEmpty()) {
      mainMemoryScopes.pop();
    }
    else {
      functionScopes.peek().pop();
    }    
  }
  
  public GraphrDataType mathOperation(String op, GraphrDataType lVal, GraphrDataType rVal) {
    switch (op.charAt(0)) {
      case '+':
        return lVal.add(rVal);
      case '-':
        return lVal.subtract(rVal);
      case '*':
        return lVal.multiply(rVal);
      case '/':
        return lVal.divide(rVal);
      case '\%':
        return lVal.mod(rVal);
    }
    return null;
  }

  public GraphrBoolean logicalOperation(String op, GraphrDataType lVal, GraphrDataType rVal) {
    if (op.equals("OR")) {
      return new GraphrBoolean(lVal.toBoolean() || rVal.toBoolean());
    }
    else if (op.equals("AND")) {
      return new GraphrBoolean(lVal.toBoolean() && rVal.toBoolean());
    }
    else if (op.equals("==")) {
      return lVal.equals(rVal);
    }
    else if (op.equals("!=")) {
      return lVal.equals(rVal).not();
    }
    else if (op.equals("<")) {
      return lVal.lessThan(rVal);
    }
    else if (op.equals("<=")) {
      return lVal.lessThanEqualTo(rVal);
    }
    else if (op.equals(">")) {
      return lVal.greaterThan(rVal);
    }
    else if (op.equals(">=")) {
      return lVal.greaterThanEqualTo(rVal);
    }
    return null;
  }
}

program 
@init {
  mainMemoryScopes.push(memory);
}
  : ( include | expr | func_def )*;

include
  : ^('include' STRING_LITERAL); // we're handling these during preproccessing. they can just be ignored here

func_def 
  : ^(FUNC name=IDENTIFIER argument* {functionMarkers.put($name.text, input.mark());} . {functionEndMarkers.put($name.text, input.mark());});

argument : ^(ARG IDENTIFIER);

block : ^(SLIST (expr)+);

expr 	: io_expr|assignment_expr|return_expr|math_expr|func_call_expr|flow_stmt;

io_expr	: ^(PUTS stmt) {System.out.println($stmt.value.toString());};

assignment_expr 
	:	assignment_stmt;

return_expr
	:	^('return' stmt) {functionReturnValues.push($stmt.value); input.seek(calledFunctionEndMarkers.pop());};

math_expr
	:	add_stmt; 

func_call_expr returns [GraphrDataType value]
	:	func_call_stmt {$value = $func_call_stmt.value;};

flow_stmt 
	: iterative_stmt
	| if_stmt
	| stmt;
	
stmt returns [GraphrDataType value]
	: constant {$value = $constant.value;}
	| array_or_hash_access {$value = $array_or_hash_access.value;}
	| IDENTIFIER {$value = getValue($IDENTIFIER.text);}
	|	array     {$value = $array.value;}
	|	hash            {$value = $hash.value;}
  | add_stmt    {$value = $add_stmt.value;}
  | built_in_functions {$value = $built_in_functions.value;}
	| enclosed_stmt {$value = $enclosed_stmt.value;};

built_in_functions returns [GraphrDataType value]
  : file_read_stmt {$value = $file_read_stmt.value;}
  | file_write_stmt {$value = $file_write_stmt.value;}
  | create_graph_stmt {$value = $create_graph_stmt.value;}
  | create_chart_stmt {$value = $create_chart_stmt.value;}
  | split_stmt {$value = $split_stmt.value;}
  | length_stmt {$value = $length_stmt.value;}
  | substring_stmt {$value = $substring_stmt.value;}
  | contains_stmt {$value = $contains_stmt.value;};

contains_stmt returns [GraphrBoolean value]
  : ^(CONTAINS associativeArray=stmt key=stmt) {
    String keyToCheck = ((GraphrString)$key.value).getString();
    $value = new GraphrBoolean(((GraphrAssociativeArray)$associativeArray.value).contains(keyToCheck));
  };

split_stmt returns [GraphrArray value]
  : ^(SPLIT string=stmt splitString=stmt) {
    GraphrArray splitStrings = new GraphrArray();
    for (String subStr : ((GraphrString)$string.value).getString().split(((GraphrString)$splitString.value).getString())) {
      splitStrings.push(new GraphrString(subStr));
    }
    $value = splitStrings;
  };

length_stmt returns [GraphrNumber value]
  : ^(LENGTH stmt) {
    GraphrNumber val = null;
    GraphrDataType stmtValue = $stmt.value;
    if (stmtValue.getClass() == GraphrString.class) {
      val = new GraphrNumber(((GraphrString)stmtValue).getString().length());
    }
    else if (stmtValue.getClass() == GraphrArray.class) {
      val = new GraphrNumber(((GraphrArray)stmtValue).getArray().size());
    }
    else if (stmtValue.getClass() == GraphrAssociativeArray.class) {
      val = new GraphrNumber(((GraphrAssociativeArray)stmtValue).getHashMap().size());      
    }
    $value = val;};

substring_stmt returns [GraphrString value]
  : ^(SUBSTRING string=stmt start=stmt end=stmt) {$value = new GraphrString(((GraphrString)$string.value).getString().substring(((GraphrNumber)$start.value).toInt(), ((GraphrNumber)$end.value).toInt()));};

enclosed_stmt returns [GraphrDataType value]
  : func_call_stmt  {$value = $func_call_stmt.value;}
  | logical_stmt {$value = $logical_stmt.value;};

create_graph_stmt returns [GraphrBoolean value]
  : ^(CREATE_GRAPH stmt) {
    GraphrAssociativeArray args = (GraphrAssociativeArray)$stmt.value;
    try {
      GraphrString type = (GraphrString)args.getValueAt("type");
      GraphrString title = (GraphrString)args.getValueAt("title");
      GraphrString xLabel = (GraphrString)args.getValueAt("xLabel");
      GraphrString yLabel = (GraphrString)args.getValueAt("yLabel");
      GraphrArray data = (GraphrArray)args.getValueAt("data");
      GraphrBoolean legend = (GraphrBoolean)args.getValueAt("legend");
      GraphrString legendLabel = (GraphrString)args.getValueAt("legendLabel");
      GraphrString file = (GraphrString)args.getValueAt("file");
      GraphrNumber width = (GraphrNumber)args.getValueAt("width");
      GraphrNumber height = (GraphrNumber)args.getValueAt("height");
      $value = GraphrGraphics.createGraph(type, title, xLabel, yLabel, legendLabel, data, legend, file, width, height);
    } catch (Exception ex) {
      System.out.println(ex.getMessage());
    }
  };

create_chart_stmt returns [GraphrBoolean value]
  : ^(CREATE_CHART stmt) {
    GraphrAssociativeArray args = (GraphrAssociativeArray)$stmt.value;    
    try {
      GraphrString type = (GraphrString)args.getValueAt("type");
      GraphrString title = (GraphrString)args.getValueAt("title");
      GraphrString xLabel = (GraphrString)args.getValueAt("xLabel");
      GraphrString yLabel = (GraphrString)args.getValueAt("yLabel");
      GraphrAssociativeArray data = (GraphrAssociativeArray)args.getValueAt("data");
      GraphrBoolean legend = (GraphrBoolean)args.getValueAt("legend");
      GraphrString file = (GraphrString)args.getValueAt("file");
      GraphrNumber width = (GraphrNumber)args.getValueAt("width");
      GraphrNumber height = (GraphrNumber)args.getValueAt("height");
      $value = GraphrGraphics.createGraph(type, title, xLabel, yLabel, data, legend, file, width, height);
    } catch (Exception ex) {
      System.out.println(ex.getMessage());
    }
  };

file_read_stmt returns [GraphrArray value]
  : ^(READ_FILE stmt) {
    GraphrArray lines = new GraphrArray();
    try {
        String fileName = ((GraphrString)$stmt.value).getString();
        BufferedReader in = new BufferedReader(new FileReader(fileName));
        String line;
        while ((line = in.readLine()) != null) {
            lines.push(new GraphrString(line));
        }
        in.close();
    } catch (IOException e) {
      System.out.println("there was an error reading the file: \n" + e.getMessage());
      $value = new GraphrArray();
    }
    $value = lines;
  };

file_write_stmt returns [GraphrBoolean value]
  : ^(WRITE_FILE fileName=stmt lines=stmt STRING_LITERAL) {
    try {
        String fName = ((GraphrString)$fileName.value).getString();
        ArrayList<GraphrDataType> fileLines = ((GraphrArray)$lines.value).getArray();
        boolean append = $STRING_LITERAL.text.equals("a") ? true : false;
        BufferedWriter out = new BufferedWriter(new FileWriter(fName, append));
        for (GraphrDataType line : fileLines) {
          out.write(((GraphrString)line).getString() + "\n");
        }
        out.close();
    } catch (IOException e) {
      System.out.println("there was an error writing the file: \n" + e.getMessage());
      $value = new GraphrBoolean(false);
    }
    $value = new GraphrBoolean(true);    
  };

func_call_stmt returns [GraphrDataType value]
@init {
  HashMap<String, GraphrDataType> thisFunctionScope = new HashMap<String, GraphrDataType>();
  Stack<Map> newScopeStack = new Stack<Map>();
  newScopeStack.push(thisFunctionScope);
  int argNumber = 1;
  GraphrWalker include = null;
  
  /* for some dumb reason the finally block gets called right after init and then again
   * after the actual rule code gets exectuted. We only want it done at the end so this
   * boolean is to ensure that. What a lame hack.
  */
  boolean inInit = true;
}
	: ^(CALL IDENTIFIER (stmt {
    CommonTree funcRoot = functions.get($IDENTIFIER.text);
    if (funcRoot != null) {
  	  String argName = funcRoot.getChild(argNumber).getChild(0).getText();
  	  argNumber++;
  	  thisFunctionScope.put(argName, $stmt.value);      
    }
    else {
      System.err.println("no defined function: "+ $IDENTIFIER.text);      
    }})*) {
      functionScopes.push(newScopeStack);
      int functionValuesSizeBeforeCall = functionReturnValues.size();
      int afterFuncCallMarker = input.mark();
      input.seek(functionMarkers.get($IDENTIFIER.text));
      calledFunctionEndMarkers.push(functionEndMarkers.get($IDENTIFIER.text));
      block();
      input.seek(afterFuncCallMarker);        
      if (functionValuesSizeBeforeCall == functionReturnValues.size()) {
        calledFunctionEndMarkers.pop();
        $value = new GraphrBoolean(true);
      }
      else {
        $value = functionReturnValues.pop();
      }
      inInit = false;
    };
finally {
  // here's that hack to make sure finally only gets called at the end
  if (!inInit) {functionScopes.pop();}
}

iterative_stmt
@init {
  int blockMarker = -1, logicalMarker = -1, incrementMarker = -1, blockEnd = -1;
  String iterativeVariableName = "";
  pushScope();
}
	:	^('for' iterative_var {iterativeVariableName = $iterative_var.variableName; logicalMarker = input.mark();} . {incrementMarker = input.mark();} . {blockMarker = input.mark();} .) {
    blockEnd = input.mark();
	  input.rewind(logicalMarker);
	  GraphrBoolean lb = (GraphrBoolean)logical_stmt();
	  while (lb.getBoolean()) {
  	  input.seek(blockMarker);
	    block();
	    input.seek(incrementMarker);
	    inc_stmt();
	    input.seek(logicalMarker);
	    lb = (GraphrBoolean)logical_stmt();
	  }
	  input.seek(blockEnd);
	}
	|	^('while' {logicalMarker = input.mark() + 1;} . {blockMarker = input.mark();} .) {
	  blockEnd = input.mark();
	  input.seek(logicalMarker);
	  GraphrBoolean whileStmt = (GraphrBoolean)logical_stmt();
    while (whileStmt.getBoolean()) {
      input.seek(blockMarker);
      block();
      input.seek(logicalMarker);
      whileStmt = (GraphrBoolean)stmt();
    }
    input.seek(blockEnd);
	}
	|	^('foreach' IDENTIFIER {iterativeVariableName = $IDENTIFIER.text; incrementMarker = input.mark();} . {blockMarker = input.mark();} .) {
	  blockEnd = input.mark();
	  input.rewind(incrementMarker);
    GraphrCollection forEachIterableItem = (GraphrCollection)stmt();
    for (GraphrDataType item : forEachIterableItem.getArray()) {
      setValue(iterativeVariableName, item);
      input.seek(blockMarker);
      block();
    }
    input.seek(blockEnd);
	}
	;
finally {
  popScope();
}
	
foreach_iterable_item returns [GraphrArray value]
  : IDENTIFIER {$value = (GraphrArray)getValue($IDENTIFIER.text);}
  | array {$value = $array.value;}
  | file_read_stmt {$value = $file_read_stmt.value;};

iterative_var returns [String variableName]
  : iterative_var_assignment {$variableName = $iterative_var_assignment.variableName;}
  | IDENTIFIER {$variableName = $IDENTIFIER.text;};

iterative_var_assignment returns [String variableName]
  : ^('=' IDENTIFIER stmt) {$variableName = $IDENTIFIER.text; setValue($variableName, $stmt.value); };

if_stmt	
@init {
  GraphrBoolean ifStmt = new GraphrBoolean(false);
  pushScope();
}
  :	^('if' logical_stmt {ifStmt = (GraphrBoolean)$logical_stmt.value; if (ifStmt.getBoolean()) {int m = input.mark(); block(); input.rewind(m);};} . ('else' {if (!ifStmt.getBoolean()) {int m = input.mark(); block(); input.rewind(m);};} .)?);
finally {
  popScope();
}

assignment_stmt	returns [GraphrDataType value]
  : ^('=' IDENTIFIER stmt) {$value = setValue($IDENTIFIER.text, $stmt.value);}
  | inc_stmt {$value = $inc_stmt.value;}
  | array_or_hash_assignment {$value = $array_or_hash_assignment.value;};

array_or_hash_assignment returns [GraphrDataType value]
  : ^('=' IDENTIFIER k=stmt v=stmt) {
    GraphrCollection collection = (GraphrCollection)getValue($IDENTIFIER.text);
    collection.setValueAt($k.value, $v.value);
    $value = $v.value;
  }
  | array_or_hash_inc_stmt {$value = $array_or_hash_inc_stmt.value;};

array_or_hash_inc_stmt returns [GraphrDataType value]
  : ^(inc_op IDENTIFIER k=stmt v=stmt) {
    GraphrCollection collection = (GraphrCollection)getValue($IDENTIFIER.text);
    GraphrDataType returnValue = mathOperation($inc_op.value, collection.getValueAt($k.value), $v.value);
    collection.setValueAt($k.value, returnValue);
    $value = returnValue;    
  };

array returns [GraphrArray value]
	: ^(ARRAY_LIST NIL) {$value = new GraphrArray();}
	| ^(ARRAY_LIST {$value = new GraphrArray();} (stmt {$value.push($stmt.value);})*) ;

array_or_hash_access returns [GraphrDataType value]
  : ^(GET_VAL IDENTIFIER stmt) {
    GraphrCollection collection = (GraphrCollection)getValue($IDENTIFIER.text);
    $value = collection.getValueAt($stmt.value); 
  };

hash returns [GraphrAssociativeArray value]
@init {
  GraphrAssociativeArray h = new GraphrAssociativeArray();
}
 	: ^(HASH_MAP (hash_element {h.put($hash_element.key, $hash_element.value);})*);
finally {
  $value = h;
}

hash_element returns [String key, GraphrDataType value]
	:	^(STRING_LITERAL constant) {$key = $STRING_LITERAL.text.substring(1, $STRING_LITERAL.text.length() - 1); $value = $constant.value;}
	| ^(STRING_LITERAL IDENTIFIER) {$key = $STRING_LITERAL.text.substring(1, $STRING_LITERAL.text.length() - 1); $value = getValue($IDENTIFIER.text);}
	| ^(STRING_LITERAL stmt) {$key = $STRING_LITERAL.text.substring(1, $STRING_LITERAL.text.length() - 1); $value = $stmt.value;};

logical_stmt returns [GraphrDataType value]
	:	^(logical_op lVal=logical_stmt rVal=logical_stmt) {$value = logicalOperation($logical_op.value, $lVal.value, $rVal.value);}
	| ^(logical_op logical_val logical_stmt) {$value = logicalOperation($logical_op.value, $logical_val.value, $logical_stmt.value);}
	| logical_val {
  	  if ($logical_val.value.getClass() == GraphrBoolean.class) {
  	    $value = $logical_val.value;
  	  }
  	  else {
  	    $value = $logical_val.value;
  	  }
  	}
	| ^(NOT logical_val) {
	  if ($logical_val.value.getClass() == GraphrBoolean.class) {
	    $value = ((GraphrBoolean)$logical_val.value).not();
	  }
	  else {
	    $value = $logical_val.value == null ? new GraphrBoolean(true) : new GraphrBoolean(false);
	  }	  
	  };

logical_val returns [GraphrDataType value]
  : func_call_stmt {$value = $func_call_stmt.value;}
  | 'true' {$value = new GraphrBoolean(true);}
  | 'false' {$value = new GraphrBoolean(false);}
  | IDENTIFIER {$value = getValue($IDENTIFIER.text);}
  | constant {$value = $constant.value;};

logical_op returns [String value]
  : or {$value = $or.value;}
  | and {$value = $and.value;}
  | '==' {$value = "==";}
  | '!=' {$value = "!=";}
  | '<' {$value = "<";}
  | '>' {$value = ">";}
  | '<=' {$value = "<=";}
  | '>=' {$value = ">=";};

or returns [String value] : OR {$value = "OR";};
and returns [String value] : AND {$value = "AND";};

inc_stmt returns [GraphrDataType value]
	:	^(inc_op IDENTIFIER stmt) {$value = setValue($IDENTIFIER.text, mathOperation($inc_op.value, getValue($IDENTIFIER.text), $stmt.value));};

inc_op returns [String value]	
  :	'+=' {$value = "+";}
  | '-=' {$value = "-";}
  | '*=' {$value = "*";}
  | '/=' {$value = "/";};

add_stmt returns [GraphrDataType value]
	: ^(add_op lVal=add_stmt rVal=add_stmt) {$value = mathOperation($add_op.value, lVal, rVal);}
	| ^(add_op lVal=mult_stmt rVal=add_stmt) {$value = mathOperation($add_op.value, lVal, rVal);}
	|	mult_stmt {value = $mult_stmt.value;};

add_op returns [String value] 
  : '+' {$value = "+";}
  | '-' {$value = "-";};

mult_stmt returns [GraphrDataType value]
	:	^(mult_op lVal=add_stmt rVal=add_stmt) {$value = mathOperation($mult_op.value, lVal, rVal);}
	| ^(mult_op lVal=mult_lval rVal=mult_stmt) {$value = mathOperation($mult_op.value, lVal, rVal);}
	| mult_lval {$value = $mult_lval.value;};

mult_lval returns [GraphrDataType value]
  : constant {$value = $constant.value;}
  | IDENTIFIER {$value = getValue($IDENTIFIER.text);}
  | func_call_stmt {$value = $func_call_stmt.value;}
  | array {$value = $array.value;};

mult_op returns [String value] 
  : '*' {$value = "*";} 
  | '/' {$value = "/";}
  | '%' {$value = "\%";};
	
constant returns [GraphrDataType value]
	:	^(HEX HEX_LITERAL)       {$value = new GraphrNumber(Integer.parseInt($HEX_LITERAL.text, 16));}
	|	^(NUMBER NUMBER_LITERAL) {$value = new GraphrNumber(Double.parseDouble($NUMBER_LITERAL.text));}
	|	^(STRING STRING_LITERAL) {$value = new GraphrString($STRING_LITERAL.text.substring(1, $STRING_LITERAL.text.length() - 1));}
	| ^(BOOL 'false')          {$value = new GraphrBoolean(new Boolean(false));}
	| ^(BOOL 'true')           {$value = new GraphrBoolean(new Boolean(true));}
	;
